﻿using System;
using System.Collections.Generic;
using System.Linq;
using Allegro.MathInterface;

namespace Allegro.Mathlib
{
    public interface IParametricResultset<T>
    {
        // <summary>
        /// Add a new set of data for parameter vector p
        /// </summary>
        /// <param name="p">Parameter values</param>
        /// <param name="data"></param>
        void AddData(double[] p, T data);

        /// <summary>
        /// Length of the resultset (number of data points added)
        /// </summary>
        int Length { get; }

        /// <summary>
        /// Array of parameter vectors. Parameters[i] is parameter vector #i, and Parameters[i][j] is
        /// the value of parameter j in parameter vector # i.
        /// </summary>
        double[][] Parameters
        {
            get;
        }
        /// <summary>
        /// Array of results data. Data[j] is the result corresponding to Parameters[j].
        /// </summary>
        T[] Data
        {
            get;
        }
    }

    public class ParametricResultset<T> : IParametricResultset<T>
    {
        private List<T> _data;
        private List<double[]> _parameters;

        public ParametricResultset()
        {
            _data = new List<T>();
            _parameters = new List<double[]>();
        }
        #region Implementation of IParametricResultset<T>

        public void AddData(double[] p, T data)
        {
            _data.Add(data);
            _parameters.Add(p);
        }

        public int Length
        {
            get { return _data.Count; }
        }

        public double[][] Parameters
        {
            get { return _parameters.ToArray(); }
        }

        public T[] Data
        {
            get { return _data.ToArray(); }
        }

        #endregion
    }
    /// <summary>
    /// MCMC, Metropolos algorithm, cf. http://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm
    /// </summary>
    public class MarkovChainMonteCarloMethod
    {
        private int _dimension;
        private ParametricResultset<double> _resultset;
        //private NormalDistribution[] _proposalDistributions;
        private NormalRandom _proposalRandom;
        private Random _random;
        private double[] _proposalSigma;
        private double[] _maxPos;
        private double[] _minPos;
        private double _maxValue;
        private double _minValue;
        private double[] _pos;
        private double[] _proposal;
        private int _nAccept;
        private int _nSteps;
        private int _nAcceptLastSteps; // Number of accepted steps last N steps

        /// <summary>
        /// Initialize a MCMC chain. nIterations is the max. number of iterations after burn-in.
        /// dimension is the number of parameters. Adjust the burn-in period by setting the BurnInDuarition property.
        /// </summary>
        /// <param name="nIterations"></param>
        /// <param name="dimension"></param>
        /// <param name="function"></param>
        public MarkovChainMonteCarloMethod(int nIterations, int dimension, VectorFunction function, double[] propsalSigma)
        {
            _nAccept = 0;
            _nSteps = 0;
            _nAcceptLastSteps = 0;

            _dimension = dimension;
            Func = function;
            NIterations = nIterations;
            _proposalSigma = propsalSigma;
            _maxPos = new double[dimension];
            _minPos = new double[dimension];
            _pos = new double[dimension];
            _proposal = new double[dimension];

            BurnInDuration = 100;
            AcceptanceRateWindow = BurnInDuration / 2;
            _resultset = new ParametricResultset<double>();

            if (dimension == 1)
                DesiredAcceptanceRate = 0.5;
            else if (dimension == 2)
                DesiredAcceptanceRate = 0.35;
            else
            {
                DesiredAcceptanceRate = 0.25;
            }

            //_proposalDistributions = new NormalDistribution[dimension];
            _proposalRandom = new NormalRandom();
            _random = new Random();
        }

        /// <summary>
        /// Max. value
        /// </summary>
        public double MaxValue
        {
            get { return _maxValue; }
        }
        /// <summary>
        /// Parameter vector of max value
        /// </summary>
        public double[] Maximum
        {
            get { return _maxPos; }
        }
        /// <summary>
        /// Min. value
        /// </summary>
        public double MinValue
        {
            get { return _minValue; }
        }
        /// <summary>
        /// Parameter vector of min value
        /// </summary>
        public double[] Minimum
        {
            get { return _minPos; }
        }
        /// <summary>
        /// Number of iterations for the burn-in period
        /// </summary>
        public int BurnInDuration
        {
            get;
            set;
        }
        /// <summary>
        /// Size of the window (#steps) over which the acceptance factor is calculated 
        /// </summary>
        public int AcceptanceRateWindow
        {
            get;
            set;
        }
        /// <summary>
        /// Current acceptance rate
        /// </summary>
        public double AcceptanceRate
        {
            get; private set;
        }
        /// <summary>
        /// Desired acceptance rate
        /// </summary>
        public double DesiredAcceptanceRate
        {
            get;
            internal set;
        }
        /// <summary>
        /// Max. number of iterations after burn-in
        /// </summary>
        public int NIterations
        {
            get;
            private set;
        }
        /// <summary>
        /// Result
        /// </summary>
        public IParametricResultset<double> Result
        {
            get { return _resultset; }
        }
        /// <summary>
        /// Vector function to be minimized
        /// </summary>
        public VectorFunction Func { get; private set; }

        /// <summary>
        /// Run a chain with p as starting point
        /// </summary>
        /// <param name="p"></param>
        /// <returns></returns>
        public double FindMaximum(double[] initialPos)
        {
            bool[] accepted = new bool[AcceptanceRateWindow];
            _maxPos = initialPos;
            _SetCurrentPos(initialPos);
            double f = _maxValue = Func(_pos);

            int nIterations = BurnInDuration + NIterations;
            for (int i = 0; i < nIterations; i++)
            {
                for (int par = 0; par < _dimension; par++)
                {
                    double x = _proposalRandom.DrawOne();
                    _proposal[par] = _proposalSigma[par] * x + _pos[par]; // New proposed pos.
                }
                double fp = Func(_proposal); // Function value at proposed pos.
                double u = _random.NextDouble();
                double r = fp / f;
                if (u <= r)
                {
                    // Accept the step
                    accepted[i % AcceptanceRateWindow] = true;
                    _nAccept++;
                    _SetCurrentPos(_proposal);
                    f = fp;
                    if(f > _maxValue)
                    {
                        _SetMaxPos(_pos);
                        _maxValue = f;
                    }
                }
                else
                {
                    accepted[i % AcceptanceRateWindow] = false; 
                }
                _nSteps++;

                if (i >= BurnInDuration)
                {
                    // We are past burn-in
                    _resultset.AddData(_pos, f);
                }
                // Compute the acceptance rate
                int nAcceptedInWindow = 0;
                for(int j = 0; j < AcceptanceRateWindow; j++)
                {
                    if (accepted[j])
                        nAcceptedInWindow++;
                }
                AcceptanceRate = ((double) nAcceptedInWindow)/AcceptanceRateWindow;
            }
            return _maxValue;
        }
        /// <summary>
        /// Run a chain with p as starting point to find the minimum of Func
        /// </summary>
        /// <param name="p"></param>
        /// <returns></returns>
        public double FindMinimum(double[] initialPos)
        {
            bool[] accepted = new bool[AcceptanceRateWindow];
            _maxPos = initialPos;
            _SetCurrentPos(initialPos);
            double f = _minValue = Func(_pos);

            int nIterations = BurnInDuration + NIterations;
            for (int i = 0; i < nIterations; i++)
            {
                for (int par = 0; par < _dimension; par++)
                {
                    double x = _proposalRandom.DrawOne();
                    _proposal[par] = _proposalSigma[par] * x + _pos[par]; // New proposed pos.
                }
                double fp = Func(_proposal); // Function value at proposed pos.
                double u = _random.NextDouble();
                double r = f / fp;
                if (u <= r)
                {
                    // Accept the step
                    accepted[i % AcceptanceRateWindow] = true;
                    _nAccept++;
                    _SetCurrentPos(_proposal);
                    f = fp;
                    if (f < _minValue)
                    {
                        _SetMinPos(_pos);
                        _minValue = f;
                    }
                }
                else
                {
                    accepted[i % AcceptanceRateWindow] = false;
                }
                _nSteps++;

                if (i >= BurnInDuration)
                {
                    // We are past burn-in
                    _resultset.AddData(_pos, f);
                }
                // Compute the acceptance rate
                int nAcceptedInWindow = 0;
                for (int j = 0; j < AcceptanceRateWindow; j++)
                {
                    if (accepted[j])
                        nAcceptedInWindow++;
                }
                AcceptanceRate = ((double)nAcceptedInWindow) / AcceptanceRateWindow;
            }
            return _minValue;
        }
        void _SetCurrentPos(double[] pos)
        {
            for (int i = 0; i < _dimension; i++)
            {
                _pos[i] = pos[i];
            }
        }
        void _SetMaxPos(double[] pos)
        {
            for (int i = 0; i < _dimension; i++)
            {
                _maxPos[i] = pos[i];
            }
        }
        void _SetMinPos(double[] pos)
        {
            for (int i = 0; i < _dimension; i++)
            {
                _minPos[i] = pos[i];
            }
        }
    }
}
